{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "\n",
    "import torch.nn.functional as F\n",
    "from typing import Any, Optional, Tuple, Type\n",
    "\n",
    "\n",
    "class PositionEmbeddingRandom(nn.Module):\n",
    "    \"\"\"\n",
    "    Positional encoding using random spatial frequencies.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None:\n",
    "        super().__init__()\n",
    "        if scale is None or scale <= 0.0:\n",
    "            scale = 1.0\n",
    "        self.register_buffer(\n",
    "            \"positional_encoding_gaussian_matrix\",\n",
    "            scale * torch.randn((2, num_pos_feats)),\n",
    "        )\n",
    "\n",
    "    def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"Positionally encode points that are normalized to [0,1].\"\"\"\n",
    "        # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape\n",
    "        coords = 2 * coords - 1\n",
    "        coords = coords @ self.positional_encoding_gaussian_matrix\n",
    "        coords = 2 * np.pi * coords\n",
    "        # outputs d_1 x ... x d_n x C shape\n",
    "        return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)\n",
    "\n",
    "    def forward(self, size: Tuple[int, int]) -> torch.Tensor:\n",
    "        \"\"\"Generate positional encoding for a grid of the specified size.\"\"\"\n",
    "        h, w = size\n",
    "        device: Any = self.positional_encoding_gaussian_matrix.device\n",
    "        grid = torch.ones((h, w), device=device, dtype=torch.float32)\n",
    "        y_embed = grid.cumsum(dim=0) - 0.5\n",
    "        x_embed = grid.cumsum(dim=1) - 0.5\n",
    "        y_embed = y_embed / h\n",
    "        x_embed = x_embed / w\n",
    "\n",
    "        pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1))\n",
    "        return pe.permute(2, 0, 1)  # C x H x W\n",
    "\n",
    "    def forward_with_coords(\n",
    "        self, coords_input: torch.Tensor, image_size: Tuple[int, int]\n",
    "    ) -> torch.Tensor:\n",
    "        \"\"\"Positionally encode points that are not normalized to [0,1].\"\"\"\n",
    "        coords = coords_input.clone()\n",
    "        coords[:, :, 0] = coords[:, :, 0] / image_size[1]\n",
    "        coords[:, :, 1] = coords[:, :, 1] / image_size[0]\n",
    "        return self._pe_encoding(coords.to(torch.float))  # B x N x C\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([128, 16, 16])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "PE = PositionEmbeddingRandom()\n",
    "PE((16, 16)).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([222, 1, 64])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sparse_embeddings = torch.empty((222, 0, 64))\n",
    "target = torch.empty((222, 1, 64))\n",
    "sparse_embeddings = torch.cat([sparse_embeddings, target], dim=1)\n",
    "sparse_embeddings.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[0, 0, 0],\n",
       "         [1, 1, 1],\n",
       "         [2, 2, 2]]),\n",
       " tensor([[0, 1, 2],\n",
       "         [0, 1, 2],\n",
       "         [0, 1, 2]]))"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x, y = torch.meshgrid(torch.arange(3), torch.arange(3))\n",
    "x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vlm",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
